#!/bin/bash
# License: GPL 
# Author: Steven Shiau <steven _at_ stevenshiau org>
# Description: Program to create a recovery ISO directly from the machine by Clonezilla live.
# To use it:
# Prepare:
# (1) A bootable Clonezilla live, either on CD or USB flash drive.
# (2) A storage space for saving the iso. Due to the temp image space, it should be double space which the image have. It's recommended to have equal space for the used blocks on the source disk.
# Then
# (1) Boot Clonezilla live. Do not choose "TO RAM" mode. Use normal mode.
# (3) Run this program. Without any option it will enter interactive mode. You can also use some options. Run with option "-h" for more info.
# Then it will take the hard drive (say /dev/sda) as an image, put the image on storage device (e.g. /dev/sdb1). Then convert that to a recovery ISO. You will have a recovery ISO for later deployment or recovery. All is done in this program.

DRBL_SCRIPT_PATH="${DRBL_SCRIPT_PATH:-/usr/share/drbl}"
. $DRBL_SCRIPT_PATH/sbin/drbl-conf-functions
. /etc/drbl/drbl-ocs.conf
. $DRBL_SCRIPT_PATH/sbin/ocs-functions

# Settings
# ocs_imgchk_batch_mode could be: on, off
ocs_imgchk_batch_mode="off"
# ocs_imgrest_batch_mode could be: on, off
ocs_imgrest_batch_mode="off"
tui="true"   # Default to turn on TUI
chk_ocsroot_mountpont="yes"
# The destination disk when restoring. By default we will ask user.
preset_dest_dev="ask_user"
#
postaction_def="choose"

#
USAGE() {
    echo "$ocs - To create a deployment iso"
    echo "Usage:"
    echo "To run $ocs:"
    echo "$ocs [OPTION] SRC_HARD_DRIVE ISO_FILE_NAME"
    echo "Options:"
    echo "-a, --postaction [choose|poweroff|reboot|command|CMD]  Set the action when restoration finishes. It can be poweroff, reboot, command (in command line prompt) or run CMD. Default is $postaction_def"
    echo "-b, --batch            Run image checking in batch mode"
    echo "-br, --batch-restore   Set image restoring as batch mode when using the created iso (Dangerous!)"
    echo "-d, --working-dir  DIR   Assign the working dir as DIR. If not assigned, $ocsroot will be used."
    echo "-nogui, --nogui    Do not show GUI (TUI) of Partclone or Partimage, use text only"
    echo "-p, --preset-dest-dev  DEV   Preset the destination device (e.g., sdg, vda...) when restoring. If not assigned, $preset_dest_dev will be used."
    echo "-s, --skip-ocsroot-mountpoint-chk   Skip checking if Clonezilla image $ocsroot is a mountpoint."
    echo "-x, --extra-boot-param  EXTRA_PARAM  Assign extra boot parameter EXTRA_PARAM for clonezilla live kernel to read. These parameters are the same with that from live-boot or live-config. Ex. \"noeject\" can be use to not prompt to eject the CD on reboot."
    echo "SRC_HARD_DRIVE is the hard drive to be taken as an image, e.g. sda, sdb..."
    echo "ISO_FILE_NAME is the name you want to give for the created file."
    echo "If no SRC_HARD_DRIVE/ISO_FILE_NAME is specified, a dialog menu will be shown."
    echo "Ex:"
    echo "To take the /dev/sda as a deployment iso, run in batch mode:"
    echo "   $ocs -b sda pragma-node.iso"
    echo
} # end of USAGE

####################
### Main program ###
####################

ocs_file="$0"
ocs=`basename $ocs_file`
#
while [ $# -gt 0 ]; do
 case "$1" in
   -b|--batch) ocs_imgchk_batch_mode="on"; shift;;
   -br|--batch-restore) ocs_imgrest_batch_mode="on"; shift;;
   -s|--skip-ocsroot-mountpoint-chk) chk_ocsroot_mountpont="no"; shift;;
   -d|--working-dir) 
            # overwrite the ocsroot in drbl.conf
            shift
            if [ -z "$(echo $1 |grep ^-.)" ]; then
              # skip the -xx option, in case 
              ocsroot="$1"
	      shift
            fi
            [ -z "$ocsroot" ] && USAGE && exit 1
	    ;;
   -p|--preset-dest-dev) 
            # Preset the device when restoring
            shift
            if [ -z "$(echo $1 |grep ^-.)" ]; then
              # skip the -xx option, in case 
              preset_dest_dev="$1"
	      shift
            fi
            [ -z "$preset_dest_dev" ] && USAGE && exit 1
	    ;;
   -nogui|--nogui)
            shift; 
            # -nogui is for backward compatable, better to use --nogui
            tui="false"
            ;;
   -x|--extra-boot-param)
            shift
            if [ -z "$(echo $1 |grep ^-.)" ]; then
              # skip the -xx option, in case 
              live_extra_boot_param="$1"
              shift
            fi
            [ -z "$live_extra_boot_param" ] && USAGE && exit 1
            ;;
   -a|--postaction) 
            # Preset the device when restoring
            shift
            if [ -z "$(echo $1 |grep ^-.)" ]; then
              # skip the -xx option, in case 
              postaction="$1"
	      shift
            fi
            [ -z "$postaction" ] && USAGE && exit 1
	    ;;
   -*)     echo "${0}: ${1}: invalid option" >&2
           USAGE >& 2
           exit 2 ;;
   *)      break ;;
 esac
done

target_hd="$1"
dest_fname="$2"

#
ask_and_load_lang_set

[ -z "$postaction" ] && postaction=$postaction_def
#
if [ -z "$target_hd" ]; then
  get_target_hd_name_from_local_machine "Choose the source disk to be imaged for the recovery" "menu"
else
  # Strip the leading /dev/ if it's assigned. Make it like sda, sdb, not /dev/sda, /dev/sdb.
  target_hd="${target_hd#/dev/*}"  # disk name, e.g. sda
fi

# Strip the leading /dev/ if it's assigned. Make it like sda, sdb, not /dev/sda, /dev/sdb.
preset_dest_dev="${preset_dest_dev#/dev/*}"  # disk name, e.g. vda

if [ -z "$dest_fname" ]; then
  # Todo: add prompt
  get_target_dir_name_when_saving  # get $target_dir
  dest_fname="$target_dir"
else
  # Strip the ".iso" in the end, because ocs-iso will add that
  dest_fname="$(echo "$dest_fname" | sed -r -e "s/.iso$//g")"
fi

#
if [ "$ocs_imgchk_batch_mode" = "on" ]; then
  ocs_imgchk_batch_opt="--batch"
else
  ocs_imgchk_batch_opt="-c"
fi
#
if [ "$ocs_imgrest_batch_mode" = "on" ]; then
  ocs_imgrest_batch_opt="--batch"
else
  ocs_imgrest_batch_opt="-c"
fi
#
if [ "$tui" = "false" ]; then
  ocs_sr_nogui_opt="--nogui"
fi

#
if [ -n "$live_extra_boot_param" ]; then
  extra_boot_param_opt="-x \"$live_extra_boot_param\""
fi

# Give warning if $ocsroot is not a mount point.
[ "$chk_ocsroot_mountpont" = "yes" ] && check_if_ocsroot_a_mountpoint


# Save the disk as an image
ocs-sr $ocs_imgchk_batch_opt $ocs_sr_nogui_opt -q2 -j2 -z1p -i 2000 -sc -p true savedisk $dest_fname $target_hd
rc_sr="$?"
if [ "$rc_sr" -eq 0 ]; then
  # Todo: Modify isolinux to add title?
  echo $msg_delimiter_star_line
  (
  cd $ocsroot
  ocs_iso_cmd="ocs-iso $extra_boot_param_opt -g en_US.UTF-8 -t -k NONE -e \"-g auto -e1 auto -e2 $ocs_imgrest_batch_opt -r -j2 -scr -k1 -p $postaction restoredisk $dest_fname $preset_dest_dev\" $dest_fname"
  echo "Running: $ocs_iso_cmd"
  eval $ocs_iso_cmd
  )
  echo "done!"
fi
